# -*- coding: utf-8 -*-

"""2698. 求一个整数的惩罚数
给你一个正整数 n ，请你返回 n 的 惩罚数 。
n 的 惩罚数 定义为所有满足以下条件 i 的数的平方和：
1 <= i <= n
i * i 的十进制表示的字符串可以分割成若干连续子字符串，且这些子字符串对应的整数值之和等于 i 。

示例 1：
输入：n = 10
输出：182
解释：总共有 3 个整数 i 满足要求：
- 1 ，因为 1 * 1 = 1
- 9 ，因为 9 * 9 = 81 ，且 81 可以分割成 8 + 1 。
- 10 ，因为 10 * 10 = 100 ，且 100 可以分割成 10 + 0 。
因此，10 的惩罚数为 1 + 81 + 100 = 182

示例 2：
输入：n = 37
输出：1478
解释：总共有 4 个整数 i 满足要求：
- 1 ，因为 1 * 1 = 1
- 9 ，因为 9 * 9 = 81 ，且 81 可以分割成 8 + 1 。
- 10 ，因为 10 * 10 = 100 ，且 100 可以分割成 10 + 0 。
- 36 ，因为 36 * 36 = 1296 ，且 1296 可以分割成 1 + 29 + 6 。
因此，37 的惩罚数为 1 + 81 + 100 + 1296 = 1478

提示：
1 <= n <= 1000"""

class Solution:
    all_punishment_number = [1, 9, 10, 36, 45, 55, 82, 91, 99, 100, 235, 297, 369, 370, 379, 414, 657, 675, 703, 756, 792, 909, 918, 945, 964, 990, 991, 999, 1000]

    @classmethod
    def solving_punishment_number(cls, upper_threshold=1000):
        def is_punishment_number(accumulate, i):
            if accumulate == i:
                return True

            x, next_accumulate, next_i = 10, accumulate, i
            while True:
                next_accumulate, next_i = accumulate // x, i - accumulate % x

                if next_accumulate == 0 or next_i < 0:
                    break

                if is_punishment_number(next_accumulate, next_i):
                    return True

                x = x*10

            return False
        
        for i in range(1, upper_threshold+1):
            if is_punishment_number(accumulate=i*i, i=i):
                cls.all_punishment_number.append(i)

    def punishmentNumber(self, n: int) -> int:
        head, tail = 0, len(self.all_punishment_number)-1

        if self.all_punishment_number[head] == n:
            p = head
        elif self.all_punishment_number[tail] == n:
            p = tail
        else:
            while True:
                middle = (head + tail) // 2
                if middle == head:
                    p = head
                    break
                if self.all_punishment_number[middle] < n:
                    head = middle
                elif self.all_punishment_number[middle] > n:
                    tail = middle
                elif self.all_punishment_number[middle] == n:
                    p = middle
                    break

        return sum([i*i for i in self.all_punishment_number[:p+1]])

if __name__ == '__main__':
    # Solution.solving_punishment_number()
    print(Solution().punishmentNumber(10))
    print(Solution().punishmentNumber(37))
    print(Solution().punishmentNumber(45))
